import tushare as ts
import pandas as pd
import time
from datetime import datetime
import schedule
import logging
from sqlalchemy import create_engine, text
import json
import os

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    filename='stock_download.log'
)

def load_config():
    """加载配置文件"""
    try:
        config_path = os.path.join(os.path.dirname(__file__), 'config.json')
        with open(config_path, 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config
    except Exception as e:
        raise Exception(f"Error loading config file: {str(e)}")

# 加载配置
CONFIG = load_config()

# 使用配置
TUSHARE_TOKEN = CONFIG['tushare']['token']
DOWNLOAD_TIME = CONFIG['download']['time']
MAX_REQUESTS_PER_MINUTE = CONFIG['download']['max_requests_per_minute']
MYSQL_HOST = CONFIG['mysql']['host']
MYSQL_PORT = CONFIG['mysql']['port']
MYSQL_USER = CONFIG['mysql']['user']
MYSQL_PASSWORD = CONFIG['mysql']['password']
MYSQL_DATABASE = CONFIG['mysql']['database']

class StockDownloader:
    def __init__(self):
        # 初始化tushare
        self.pro = ts.pro_api(TUSHARE_TOKEN)
        self.request_count = 0
        self.last_request_time = time.time()
        
        # 初始化数据库
        self._init_database()
    
    def _init_database(self):
        """初始化数据库和表"""
        try:
            #日志打印mysql连接信息
            logging.info(f"Connecting to MySQL server at {MYSQL_HOST}:{MYSQL_PORT},user:{MYSQL_USER},password:{MYSQL_PASSWORD},database:{MYSQL_DATABASE }")
            # 首先创建与MySQL服务器的连接（不指定数据库）
            engine_server = create_engine(
                f'mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}'
            )
            
            # 创建数据库（如果不存在）
            with engine_server.connect() as conn:
                conn.execute(text(f"CREATE DATABASE IF NOT EXISTS {MYSQL_DATABASE}"))
                conn.commit()  # 确保更改被提交
            
            # 创建指向特定数据库的连接
            self.engine = create_engine(
                f'mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}'
            )
            
            # 创建表（如果不存在）
            create_bardata_table = """
            CREATE TABLE IF NOT EXISTS `dbbardata` (
                `id` int NOT NULL AUTO_INCREMENT,
                `symbol` varchar(255) COLLATE utf8mb4_general_ci NOT NULL,
                `exchange` varchar(255) COLLATE utf8mb4_general_ci NOT NULL,
                `datetime` datetime NOT NULL,
                `interval` varchar(255) COLLATE utf8mb4_general_ci NOT NULL,
                `volume` double NOT NULL,
                `turnover` double NOT NULL,
                `open_interest` double NOT NULL,
                `open_price` double NOT NULL,
                `high_price` double NOT NULL,
                `low_price` double NOT NULL,
                `close_price` double NOT NULL,
                PRIMARY KEY (`id`),
                UNIQUE KEY `dbbardata_symbol_exchange_interval_datetime` 
                    (`symbol`, `exchange`, `interval`, `datetime`)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci;
            """
            
            create_overview_table = """
            CREATE TABLE IF NOT EXISTS `dbbaroverview` (
                `id` int NOT NULL AUTO_INCREMENT,
                `symbol` varchar(255) COLLATE utf8mb4_general_ci NOT NULL,
                `exchange` varchar(255) COLLATE utf8mb4_general_ci NOT NULL,
                `interval` varchar(255) COLLATE utf8mb4_general_ci NOT NULL,
                `count` int NOT NULL,
                `start` datetime NOT NULL,
                `end` datetime NOT NULL,
                PRIMARY KEY (`id`),
                UNIQUE KEY `dbbaroverview_symbol_exchange_interval` 
                    (`symbol`, `exchange`, `interval`)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci;
            """
            
            create_stock_basic_table = """
            CREATE TABLE IF NOT EXISTS `stock_basic` (
                `id` int NOT NULL AUTO_INCREMENT,
                `ts_code` varchar(20) COLLATE utf8mb4_general_ci NOT NULL COMMENT 'TS代码',
                `symbol` varchar(20) COLLATE utf8mb4_general_ci NOT NULL COMMENT '股票代码',
                `name` varchar(100) COLLATE utf8mb4_general_ci NOT NULL COMMENT '股票名称',
                `area` varchar(50) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '地域',
                `industry` varchar(50) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '所属行业',
                `fullname` varchar(200) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '股票全称',
                `enname` varchar(200) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '英文全称',
                `cnspell` varchar(50) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '拼音缩写',
                `market` varchar(50) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '市场类型',
                `exchange` varchar(20) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '交易所代码',
                `curr_type` varchar(20) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '交易货币',
                `list_status` char(1) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '上市状态',
                `list_date` date DEFAULT NULL COMMENT '上市日期',
                `delist_date` date DEFAULT NULL COMMENT '退市日期',
                `is_hs` char(1) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '是否沪深港通标的',
                `act_name` varchar(100) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '实控人名称',
                `act_ent_type` varchar(50) COLLATE utf8mb4_general_ci DEFAULT NULL COMMENT '实控人企业性质',
                `updated_at` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP COMMENT '更新时间',
                PRIMARY KEY (`id`),
                UNIQUE KEY `uk_ts_code` (`ts_code`),
                UNIQUE KEY `uk_symbol` (`symbol`),
                KEY `idx_list_status` (`list_status`),
                KEY `idx_market` (`market`),
                KEY `idx_industry` (`industry`),
                KEY `idx_area` (`area`)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci COMMENT='股票基本信息表';
            """
            
            # 执行建表语句
            with self.engine.connect() as conn:
                conn.execute(text(create_bardata_table))
                conn.execute(text(create_overview_table))
                conn.execute(text(create_stock_basic_table))
                conn.commit()  # 确保更改被提交
            
            logging.info("Database and tables initialization completed successfully")
            
        except Exception as e:
            logging.error(f"Error initializing database: {str(e)}")
            raise
    
    def _check_rate_limit(self):
        """控制请求频率"""
        current_time = time.time()
        if current_time - self.last_request_time >= 60:
            self.request_count = 0
            self.last_request_time = current_time
        
        if self.request_count >= MAX_REQUESTS_PER_MINUTE:
            sleep_time = 60 - (current_time - self.last_request_time)
            if sleep_time > 0:
                time.sleep(sleep_time)
            self.request_count = 0
            self.last_request_time = time.time()
    
    def _process_data(self, df, symbol):
        """处理数据格式"""
        # 重命名列以匹配数据库结构
        df = df.rename(columns={
            'ts_code': 'symbol',
            'trade_date': 'datetime',
            'vol': 'volume',
            'amount': 'turnover',
            'open': 'open_price',
            'high': 'high_price',
            'low': 'low_price',
            'close': 'close_price'
        })
        
        # 添加固定列
        df['exchange'] = 'SSE' if symbol.endswith('.SH') else 'SZSE'
        df['interval'] = 'd'
        df['open_interest'] = 0  # 期货相关，股票设为0
        
        # 处理symbol，去除.SH和.SZ后缀
        df['symbol'] = df['symbol'].str.replace(r'\.(SH|SZ)$', '', regex=True)
        
        # 转换日期格式
        df['datetime'] = pd.to_datetime(df['datetime'], format='%Y%m%d')
        
        # 选择需要的列
        columns = [
            'symbol', 'exchange', 'datetime', 'interval', 'volume', 
            'turnover', 'open_interest', 'open_price', 'high_price', 
            'low_price', 'close_price'
        ]
        
        return df[columns]

    def _update_overview(self, symbol, exchange, interval):
        """更新dbbaroverview表"""
        try:
            # 查询该品种的统计信息
            query = f"""
                SELECT 
                    COUNT(*) as `count`,
                    MIN(`datetime`) as `start`,
                    MAX(`datetime`) as `end`
                FROM `dbbardata`
                WHERE `symbol` = '{symbol}'
                AND `exchange` = '{exchange}'
                AND `interval` = '{interval}'
            """
            
            stats = pd.read_sql(query, self.engine)
            
            if not stats.empty and stats['count'].iloc[0] > 0:
                # 准备更新数据，将日期对象转换为字符串
                overview_data = {
                    'symbol': symbol,
                    'exchange': exchange,
                    'interval': interval,
                    'count': int(stats['count'].iloc[0]),
                    'start': stats['start'].iloc[0].strftime('%Y-%m-%d %H:%M:%S'),
                    'end': stats['end'].iloc[0].strftime('%Y-%m-%d %H:%M:%S')
                }
                
                # 使用REPLACE INTO来更新或插入数据
                insert_query = """
                    REPLACE INTO `dbbaroverview` 
                    (`symbol`, `exchange`, `interval`, `count`, `start`, `end`)
                    VALUES 
                    (:symbol, :exchange, :interval, :count, :start, :end)
                """
                
                with self.engine.connect() as conn:
                    conn.execute(
                        text(insert_query),
                        overview_data
                    )
                    conn.commit()
                
                logging.info(f"Successfully updated overview for {symbol}")
            else:
                logging.info(f"No data found for {symbol}.{exchange}")
                
        except Exception as e:
            logging.error(f"Error updating overview for {symbol}: {str(e)}")

    def _get_last_trade_date(self, symbol, exchange, interval):
        """获取数据库中该股票最新的交易日期"""
        query = f"""
            SELECT MAX(`datetime`) as `last_date`
            FROM `dbbardata`
            WHERE `symbol` = '{symbol}'
            AND `exchange` = '{exchange}'
            AND `interval` = '{interval}'
        """
        
        try:
            df = pd.read_sql(query, self.engine)
            if not df.empty and df['last_date'].iloc[0] is not None:
                return df['last_date'].iloc[0].strftime('%Y%m%d')
            return None
        except Exception as e:
            logging.error(f"Error getting last trade date for {symbol}: {str(e)}")
            return None

    def download_daily_data(self):
        """下载每日股票数据"""
        try:
            # 获取当前日期并转换为字符串格式
            today = datetime.now().strftime('%Y%m%d')
            
            # 获取所有股票列表
            stock_list = self.pro.stock_basic(exchange='', list_status='L', fields='ts_code,symbol,name,area,industry,fullname,enname,cnspell,market,exchange,curr_type,list_status,list_date,delist_date,is_hs,act_name,act_ent_type')
            
            # 将stock_list更新到表stock_basic
            stock_list.to_sql(
                'stock_basic',
                self.engine,
                if_exists='replace', 
                index=False
            )
        except Exception as e:
            logging.error(f"Error in get stock_basic: {str(e)}") 
            logging.error(f"load db stock_basic ")
            # 从数据库中获取stock_basic
            stock_list = pd.read_sql('select * from stock_basic', self.engine)
        
        try:
            for _, row in stock_list.iterrows():
                self._check_rate_limit()
                
                try:
                    # 获取该股票最新的数据日期
                    exchange = 'SSE' if row['ts_code'].endswith('.SH') else 'SZSE'
                    last_date = self._get_last_trade_date(row['ts_code'], exchange, 'd')
                    # 上市日期
                    list_date = row['list_date']
                    # 如果有历史数据，则从最新日期后一天开始下载
                    if last_date:
                        start_date = pd.to_datetime(last_date).strftime('%Y%m%d')
                    else:
                        # 如果没有历史数据，从上市至今的所有历史数据
                        start_date = list_date
                    
                    # 将start_date转换为datetime对象进行比较
                    start_date_dt = pd.to_datetime(start_date)
                    today_dt = pd.to_datetime(today)
                    
                    # 如果最新日期等于今天，说明数据已是最新
                    if start_date_dt >= today_dt:
                        logging.info(f"Data already up to date for {row['ts_code']}")
                        continue
                    
                    # 获取日线数据
                    df = self.pro.daily(
                        ts_code=row['ts_code'],
                        start_date=start_date,
                        end_date=today
                    )
                    
                    if df is not None and not df.empty:
                        # 处理数据
                        processed_df = self._process_data(df, row['ts_code'])
                        
                        # 写入数据库
                        processed_df.to_sql(
                            'dbbardata',
                            self.engine,
                            if_exists='append',
                            index=False,
                            chunksize=1000
                        )
                        
                        # 更新overview表
                        symbol = row['ts_code'].split('.')[0]
                        self._update_overview(symbol, exchange, 'd')
                        
                        logging.info(
                            f"Successfully updated data for {row['ts_code']} "
                            f"from {start_date} to {today}"
                        )
                        self.request_count += 1
                        if self.request_count >= 10:
                            break
                    else:
                        logging.info(f"No new data available for {row['ts_code']}")
                    
                except Exception as e:
                    logging.error(f"Error processing {row['ts_code']}: {str(e)}")
                    continue
        except Exception as e:
            logging.error(f"Error in download process: {str(e)}")           
       

def run_download():
    """执行下载任务"""
    logging.info("Starting daily download task")
    downloader = StockDownloader()
    downloader.download_daily_data()
    logging.info("Daily download task completed")

def main():
    run_download()
    # # 设置定时任务
    # schedule.every().day.at(DOWNLOAD_TIME).do(run_download)
    
    # logging.info(f"Scheduler started, will run daily at {DOWNLOAD_TIME}")
    
    # # 运行定时任务
    # while True:
    #     schedule.run_pending()
    #     time.sleep(60)

if __name__ == "__main__":
    main() 